#!/usr/bin/python3
# -*- coding: UTF-8 -*-

import rospy
from sensor_msgs.msg import Imu

class ImuCalibrator:
    def __init__(self):
        rospy.init_node('imu_calibrator')

        # 订阅和发布
        self.sub = rospy.Subscriber('/imu_raw', Imu, self.imu_callback)
        self.pub = rospy.Publisher('/imu_filter', Imu, queue_size=10)

        # 校准参数
        self.start_time = None
        self.calibrating = True
        self.gyro_sum = [0.0, 0.0, 0.0]
        self.accel_sum = [0.0, 0.0, 0.0]
        self.num_samples = 0

        # 零偏
        self.gyro_bias = [0.0, 0.0, 0.0]
        self.accel_bias = [0.0, 0.0, 0.0]

    def imu_callback(self, msg):
        if self.calibrating:
            if self.start_time is None:
                # 收到第一条消息时开始计时
                self.start_time = rospy.Time.now()
                self._accumulate_data(msg)
                rospy.loginfo("Calibration started.")
            else:
                elapsed = (rospy.Time.now() - self.start_time).to_sec()
                if elapsed < 5.0:
                    self._accumulate_data(msg)
                else:
                    # 计算零偏并结束校准
                    self._compute_bias()
                    self.calibrating = False
                    rospy.loginfo("Calibration completed.")
        else:
            # 发布校正后的数据
            self._publish_corrected(msg)

    def _accumulate_data(self, msg):
        """累加校准数据"""
        self.gyro_sum[0] += msg.angular_velocity.x
        self.gyro_sum[1] += msg.angular_velocity.y
        self.gyro_sum[2] += msg.angular_velocity.z
        self.accel_sum[0] += msg.linear_acceleration.x
        self.accel_sum[1] += msg.linear_acceleration.y
        self.accel_sum[2] += msg.linear_acceleration.z
        self.num_samples += 1

    def _compute_bias(self):
        """计算零偏"""
        if self.num_samples == 0:
            rospy.logwarn("No samples for calibration. Bias set to zero.")
            return
        self.gyro_bias = [s / self.num_samples for s in self.gyro_sum]
        self.accel_bias = [s / self.num_samples for s in self.accel_sum]
        rospy.loginfo("Gyro Bias: x=%.6f, y=%.6f, z=%.6f", *self.gyro_bias)
        rospy.loginfo("Accel Bias: x=%.6f, y=%.6f, z=%.6f", *self.accel_bias)

    def _publish_corrected(self, msg):
        """发布校正后的IMU数据"""
        corrected = Imu()
        corrected.header = msg.header

        # 校正角速度
        corrected.angular_velocity.x = msg.angular_velocity.x - self.gyro_bias[0]
        corrected.angular_velocity.y = msg.angular_velocity.y - self.gyro_bias[1]
        corrected.angular_velocity.z = msg.angular_velocity.z - self.gyro_bias[2]

        # 校正加速度
        corrected.linear_acceleration.x = msg.linear_acceleration.x - self.accel_bias[0]
        corrected.linear_acceleration.y = msg.linear_acceleration.y - self.accel_bias[1]
        corrected.linear_acceleration.z = msg.linear_acceleration.z
        # corrected.linear_acceleration.z = msg.linear_acceleration.z - self.accel_bias[2]

        # 保留原始方向及协方差信息
        corrected.orientation = msg.orientation
        corrected.angular_velocity_covariance = msg.angular_velocity_covariance
        corrected.linear_acceleration_covariance = msg.linear_acceleration_covariance
        corrected.orientation_covariance = msg.orientation_covariance

        self.pub.publish(corrected)

if __name__ == '__main__':
    try:
        calibrator = ImuCalibrator()
        rospy.spin()
    except rospy.ROSInterruptException:
        pass